import pandas as pd
import numpy as np
import cv2
from random import shuffle
npdata=np.load('train_data.npy',allow_pickle=True)
data=pd.DataFrame(npdata)

left=[]
right=[]
attack=[]
jump=[]
rush=[]
nochoice=[]
shuffle(npdata)
for each in npdata:
    img = each[0]
    choice = each[1]
    if choice == [1,0,0,0,0,0]:
        left.append([img,choice])
    elif choice == [0,1,0,0,0,0]:
        right.append([img,choice])
    elif choice == [0,0,1,0,0,0]:
        attack.append([img,choice])
    elif choice == [0,0,0,1,0,0]:
        jump.append([img,choice])
    elif choice == [0,0,0,0,1,0]:
        rush.append([img,choice])
    elif choice == [0,0,0,0,0,1]:
        nochoice.append([img,choice])
    else:
        print('error!')


nochoice=nochoice[:len((right))][:len(left)][:len(attack)][:len(jump)][:len(rush)]
right=right[:len(nochoice)]
left=left[:len(nochoice)]
attack=attack[:len(nochoice)]
jump=jump[:len(nochoice)]
rush=rush[:len(nochoice)]
print(len(right))

shuffle_data=rush+right+left+jump+attack+nochoice
shuffle(shuffle_data)
np.save('shuffle_datamin.npy',shuffle_data)
print(len(shuffle_data))
"""for each in npdata:
    img = each[0]
    choice = each[1]
    cv2.imshow('test',img)
    print(choice)
    if cv2.waitKey(25) & 0xff == ord('q'):
        cv2.destroyAllWindows()
        break"""